import unittest

import albumentations.augmentations.functional as A
import numpy as np
import torch
import torchvision.transforms.functional as B


class MyTestCase(unittest.TestCase):

    def test_normalize(self):
        mean = (0.485, 0.456, 0.406)
        std = (0.229, 0.224, 0.225)
        image = np.array([
            [
                [0, 0, 0],
                [1, 1, 1],
                [2, 2, 2],
            ],
            [
                [126, 126, 126],
                [127, 127, 127],
                [128, 128, 128],
            ],
            [
                [253, 253, 253],
                [254, 254, 254],
                [255, 255, 255]
            ]
        ])
        image = A.normalize(image, mean, std)
        print(image)

    def test_normalize2(self):
        mean = (0.485, 0.456, 0.406)
        std = (0.229, 0.224, 0.225)
        image = np.array([
            [
                [1, 1, 1],
                [2, 2, 2],
                [3, 3, 3],
            ],
            [
                [126, 126, 126],
                [127, 127, 127],
                [128, 128, 128],
            ],
            [
                [253, 253, 253],
                [254, 254, 254],
                [255, 255, 255]
            ]
        ])
        image = torch.from_numpy(image)
        image = B.normalize(image, mean, std).numpy()
        print(image)


if __name__ == '__main__':
    unittest.main()
